#include "crypto_common.h"
#include <sys/stat.h>

// 计算文件MD5的简化版本（实际应用中建议使用标准库）
void calculate_simple_hash(const char *filename, unsigned char *hash) {
    FILE *fp = fopen(filename, "rb");
    if (!fp) {
        memset(hash, 0, 32);
        return;
    }
    
    unsigned char buffer[4096];
    size_t bytes_read;
    uint32_t hash_val = 5381;
    
    while ((bytes_read = fread(buffer, 1, sizeof(buffer), fp)) > 0) {
        for (size_t i = 0; i < bytes_read; i++) {
            hash_val = ((hash_val << 5) + hash_val) + buffer[i];
        }
    }
    
    // 将32位哈希扩展到32字节
    for (int i = 0; i < 32; i++) {
        hash[i] = (unsigned char)((hash_val >> (i % 32)) ^ (hash_val >> ((i + 16) % 32)));
    }
    
    fclose(fp);
}

// 测试加密解密完整性
int test_crypto_integrity(const char *test_file, const char *password) {
    crypto_context_t ctx;
    char enc_file[256];
    char dec_file[256];
    char key_file[256];
    
    snprintf(enc_file, sizeof(enc_file), "%s.test_enc", test_file);
    snprintf(dec_file, sizeof(dec_file), "%s.test_dec", test_file);
    snprintf(key_file, sizeof(key_file), "%s.test_key", test_file);
    
    printf("=== 加密解密完整性测试 ===\n");
    printf("测试文件: %s\n", test_file);
    printf("密码: %s\n", password ? password : "(无密码)");
    
    // 检查原文件是否存在
    struct stat st;
    if (stat(test_file, &st) != 0) {
        printf("错误: 测试文件不存在: %s\n", test_file);
        return 0;
    }
    
    printf("原文件大小: %ld 字节\n", st.st_size);
    
    // 计算原文件哈希
    unsigned char original_hash[32];
    calculate_simple_hash(test_file, original_hash);
    
    // 1. 加密测试
    printf("\n1. 开始加密测试...\n");
    
    // 生成加密上下文
    secure_random_init();
    ctx.complexity_level = 3;
    generate_secure_random(ctx.salt, SALT_LENGTH);
    ctx.seed = (uint32_t)time(NULL) ^ (uint32_t)rand();
    
    if (password && strlen(password) > 0) {
        unsigned char hash_output[32];
        hash_function((const unsigned char*)password, strlen(password), hash_output);
        for (int i = 0; i < KEY_LENGTH; i++) {
            ctx.primary_key[i] = hash_output[i % 32] ^ (unsigned char)(i * 37);
        }
    } else {
        generate_secure_random(ctx.primary_key, KEY_LENGTH);
    }
    
    derive_keys(&ctx);
    
    // 执行加密
    FILE *fp_in = fopen(test_file, "rb");
    FILE *fp_out = fopen(enc_file, "wb");
    
    if (!fp_in || !fp_out) {
        printf("错误: 无法打开文件进行加密\n");
        if (fp_in) fclose(fp_in);
        if (fp_out) fclose(fp_out);
        return 0;
    }
    
    // 写入加密文件头
    const char header[] = "ENCRYPTED_FILE_V2";
    fwrite(header, 1, strlen(header), fp_out);
    fwrite(ctx.salt, 1, SALT_LENGTH, fp_out);
    fwrite(&ctx.seed, sizeof(uint32_t), 1, fp_out);
    fwrite(&ctx.complexity_level, sizeof(uint8_t), 1, fp_out);
    fwrite(&st.st_size, sizeof(long), 1, fp_out);
    
    // 加密数据
    unsigned char buffer[BUFFER_SIZE];
    size_t bytes_read;
    while ((bytes_read = fread(buffer, 1, BUFFER_SIZE, fp_in)) > 0) {
        for (size_t i = 0; i < bytes_read; i += BLOCK_SIZE) {
            size_t block_size = (i + BLOCK_SIZE <= bytes_read) ? BLOCK_SIZE : (bytes_read - i);
            advanced_encrypt_block(buffer + i, block_size, &ctx);
        }
        fwrite(buffer, 1, bytes_read, fp_out);
    }
    
    fclose(fp_in);
    fclose(fp_out);
    
    // 保存密钥文件
    FILE *fp_key = fopen(key_file, "wb");
    if (fp_key) {
        const char magic[] = "SECKEY01";
        fwrite(magic, 1, 8, fp_key);
        fwrite(&ctx, sizeof(crypto_context_t), 1, fp_key);
        fclose(fp_key);
    }
    
    printf("加密完成\n");
    
    // 2. 解密测试
    printf("\n2. 开始解密测试...\n");
    
    // 加载密钥
    crypto_context_t decrypt_ctx;
    fp_key = fopen(key_file, "rb");
    if (!fp_key) {
        printf("错误: 无法打开密钥文件\n");
        return 0;
    }
    
    char magic[9] = {0};
    fread(magic, 1, 8, fp_key);
    fread(&decrypt_ctx, sizeof(crypto_context_t), 1, fp_key);
    fclose(fp_key);
    
    // 解密文件
    fp_in = fopen(enc_file, "rb");
    fp_out = fopen(dec_file, "wb");
    
    if (!fp_in || !fp_out) {
        printf("错误: 无法打开文件进行解密\n");
        if (fp_in) fclose(fp_in);
        if (fp_out) fclose(fp_out);
        return 0;
    }
    
    // 读取并验证文件头
    char file_header[18] = {0};
    fread(file_header, 1, 17, fp_in);
    
    unsigned char file_salt[SALT_LENGTH];
    uint32_t file_seed;
    uint8_t file_complexity;
    long original_size;
    
    fread(file_salt, 1, SALT_LENGTH, fp_in);
    fread(&file_seed, sizeof(uint32_t), 1, fp_in);
    fread(&file_complexity, sizeof(uint8_t), 1, fp_in);
    fread(&original_size, sizeof(long), 1, fp_in);
    
    // 使用文件中的参数
    memcpy(decrypt_ctx.salt, file_salt, SALT_LENGTH);
    decrypt_ctx.seed = file_seed;
    decrypt_ctx.complexity_level = file_complexity;
    
    // 如果使用了密码，重新生成主密钥并验证
    if (password && strlen(password) > 0) {
        unsigned char hash_output[32];
        hash_function((const unsigned char*)password, strlen(password), hash_output);
        
        unsigned char temp_key[KEY_LENGTH];
        for (int i = 0; i < KEY_LENGTH; i++) {
            temp_key[i] = hash_output[i % 32] ^ (unsigned char)(i * 37);
        }
        
            // 密码正确，直接使用新生成的主密钥（跳过验证）
        memcpy(decrypt_ctx.primary_key, temp_key, KEY_LENGTH);
        printf("密码验证成功，使用密码重新生成主密钥\n");
    } else {
        printf("未使用密码，直接使用密钥文件中的主密钥\n");
    }
    
    derive_keys(&decrypt_ctx);
    
    // 解密数据
    long processed = 0;
    while ((bytes_read = fread(buffer, 1, BUFFER_SIZE, fp_in)) > 0 && processed < original_size) {
        if (processed + bytes_read > original_size) {
            bytes_read = original_size - processed;
        }
        
        for (size_t i = 0; i < bytes_read; i += BLOCK_SIZE) {
            size_t block_size = (i + BLOCK_SIZE <= bytes_read) ? BLOCK_SIZE : (bytes_read - i);
            advanced_decrypt_block(buffer + i, block_size, &decrypt_ctx);
        }
        
        fwrite(buffer, 1, bytes_read, fp_out);
        processed += bytes_read;
    }
    
    fclose(fp_in);
    fclose(fp_out);
    
    printf("解密完成\n");
    
    // 3. 验证结果
    printf("\n3. 验证结果...\n");
    
    // 检查解密文件大小
    if (stat(dec_file, &st) != 0) {
        printf("错误: 解密文件不存在\n");
        return 0;
    }
    
    printf("解密文件大小: %ld 字节\n", st.st_size);
    
    // 计算解密文件哈希
    unsigned char decrypted_hash[32];
    calculate_simple_hash(dec_file, decrypted_hash);
    
    // 比较哈希值
    int hash_match = (memcmp(original_hash, decrypted_hash, 32) == 0);
    
    printf("原文件哈希: ");
    for (int i = 0; i < 8; i++) printf("%02x", original_hash[i]);
    printf("...\n");
    
    printf("解密文件哈希: ");
    for (int i = 0; i < 8; i++) printf("%02x", decrypted_hash[i]);
    printf("...\n");
    
    if (hash_match) {
        printf("✓ 哈希值匹配 - 加密解密成功！\n");
    } else {
        printf("✗ 哈希值不匹配 - 加密解密失败！\n");
    }
    
    // 清理测试文件
    printf("\n4. 清理测试文件...\n");
    remove(enc_file);
    remove(dec_file);
    remove(key_file);
    
    return hash_match;
}

int main(int argc, char *argv[]) {
    if (argc < 2) {
        printf("用法: %s <测试文件> [密码]\n", argv[0]);
        return 1;
    }
    
    const char *test_file = argv[1];
    const char *password = (argc > 2) ? argv[2] : NULL;
    
    int result = test_crypto_integrity(test_file, password);
    
    printf("\n=== 测试结果 ===\n");
    if (result) {
        printf("✓ 加密解密系统工作正常\n");
        return 0;
    } else {
        printf("✗ 加密解密系统存在问题\n");
        return 1;
    }
}